#!/usr/bin/env python

import cv2
import message_filters
import rospy
from cv_bridge import CvBridge

from sensor_msgs.msg import Image

import numpy as np
robot_namespace = 'sdnyc_webots'

#import pcl_msgs
#import pcl
#import pcl.pcl_visualization

fx = 483.1147
fy = 483.1147
cx = 200
cy = 200
b = 0.12

class StereoMatch:
    def __init__(self):
        pass

    def match(self, left_image, right_image):
        window_size = 15
        min_disp = 0
        num_disp = 32

        stereo = cv2.StereoSGBM_create(minDisparity=min_disp,
                                       numDisparities=num_disp,
                                       P1=8 * 3 * window_size,
                                       P2=32 * 3 * window_size,
                                       disp12MaxDiff=1,
                                       blockSize=11,
                                       uniquenessRatio=24,
                                       speckleWindowSize=100)
        disparity = stereo.compute(left_image, right_image).astype(
            np.float32) / 16.0
        return disparity

    def pointcloud(self, disparity, colors):
        Q = np.array([[1, 0, 0, -cx], [0, 1, 0, -cy], [0, 0, 0, fx],
                      [0, 0, -1 / b, 0]])
        points_3d = cv2.reprojectImageTo3D(disparity, Q)
        rows, cols = points_3d.shape[0:2]
        size = rows * cols
        points_ = self.hw3ToN3(points_3d).astype(np.int16)
        colors_ = self.hw3ToN3(colors).astype(np.int64)
        blue = colors_[:, 0].reshape(size, 1)
        green = colors_[:, 1].reshape(size, 1)
        red = colors_[:, 2].reshape(size, 1)
        rgb = np.left_shift(blue, 0) + np.left_shift(green, 8) + np.left_shift(
            red, 16)
        pointcloud = np.hstack((points_, rgb)).astype(np.float32)
        return pointcloud

    def hw3ToN3(self, points):
        height, width = points.shape[0:2]
        points_1 = points[:, :, 0].reshape(height * width, 1)
        points_2 = points[:, :, 1].reshape(height * width, 1)
        points_3 = points[:, :, 2].reshape(height * width, 1)
        points_ = np.hstack((points_1, points_2, points_3))
        return points_

    def pointcloud_compute(self, image, disparity):
        pass
        #for v in range(disparity.shape[0]):
        #for u in range(disparity.shape[1]):
        #if disparity[v,u]<=0.0 or disparity[v,u]>=96.0:
        #continue
        #point_4d=(0,0,0,image[v,u]/255.0)
        #x=(u-cx)/fx
        #y=(v-cy)/fy
        #depth = fx*b/(disparity[v, u]
        #point[0] = x * depth
        #point[1] = y * depth
        #point[2] = depth

        #color=image.reshape(-1,3)
        #disp=disparity.reshape(-1)
        #mask = ((disp > disp.min()) &
        #np.all(~np.isnan(points), axis=1) &
        #np.all(~np.isinf(points), axis=1))
        #return points[mask], color[mask]

    def result_to_diaprity(self, result):
        pass


class Matcher:
    def __init__(self):
        self.__stereo_matcher = StereoMatch()
        self.__bridge = CvBridge()
        self.__disparity_pub = rospy.Publisher('/image_matcher/disparity',
                                               Image,
                                               queue_size=1)

    def __image_callback(self, sub_left, sub_right):
        #cv_image = self.__bridge.imgmsg_to_cv2(cam_image, 'bgr8')
        cv_l_image = self.__bridge.imgmsg_to_cv2(sub_left, 'bgr8')
        cv_r_image = self.__bridge.imgmsg_to_cv2(sub_right, 'bgr8')
        cv_left_image = cv2.cvtColor(cv_l_image, cv2.COLOR_BGR2GRAY)
        cv_right_image = cv2.cvtColor(cv_r_image, cv2.COLOR_BGR2GRAY)
        width = cv_left_image.shape[0]
        heigth = cv_left_image.shape[1]
        image = cv_l_image
        #cv2.imwrite("/home/liu/left.png",cv_left_image)
        disparitys = self.__stereo_matcher.match(cv_left_image, cv_right_image)
        #pointclouds=self.__stereo_matcher.pointcloud(disparitys,image)
        #cv2.imshow("disparity",disparitys)
        #cv2.waitKey(200)
        #self.view_cloud(pointclouds)
        msg_disparitys = self.__bridge.cv2_to_imgmsg(disparitys,
                                                     encoding='32FC1')
        self.__disparity_pub.publish(msg_disparitys)

    # def view_cloud(self, pointcloud):
    #     cloud = pcl.PointCloud_PointXYZRGBA()
    #     cloud.from_array(pointcloud)
    #     try:
    #         visual = pcl.pcl_visualization.CloudViewing()
    #         visual.ShowColorACloud(cloud)
    #         v = True
    #         while v:
    #             v = not (visual.WasStopped())
    #     except:
    #         pass

    def Run(self):
        sub_left = message_filters.Subscriber(
            '/sdnyc_webots/left_camera/image', Image)
        rospy.loginfo('stereo match left image loaded.')
        sub_right = message_filters.Subscriber(
            '/sdnyc_webots/right_camera/image', Image)
        rospy.loginfo('stereo match right image loaded.')
        ts = message_filters.ApproximateTimeSynchronizer([sub_left, sub_right],
                                                         1, 1)
        ts.registerCallback(self.__image_callback)

        rospy.spin()


if __name__ == '__main__':
    rospy.init_node('stereo_match', anonymous=False)

    stereo_matcher = Matcher()

    stereo_matcher.Run()
